import pyarrow as pa
import pyarrow.fs
import unittest


class TestAccessArrow(unittest.TestCase):

    def setUp(self) -> None:
        self.file = 'test.arrow'
        self.names = ['name', 'content']
        self.data = [
            pa.array(['name']),
            pa.array(['content'])
        ]
        self.batch = pa.record_batch(self.data, self.names)
        self.local = pyarrow.fs.LocalFileSystem()

    def test_write_arrow(self):
        with self.local.open_append_stream(self.file) as f:
            with pa.RecordBatchFileWriter(f, self.batch.schema) as writer:
                for i in range(50000):
                    writer.write_batch(self.batch)

    def test_write_arrow2(self):
        with pa.OSFile('test.arrow', 'wb') as sink:
            with pa.ipc.new_file(sink, self.batch.schema) as writer:
                writer.write(self.batch)

    def test_read_arrow(self):
        with pa.memory_map('test.arrow', 'rb') as source:
            loaded_array = pa.ipc.open_file(source).read_all()
            for item in loaded_array:
                print(item)
